import json
import requests
from requests.adapters import HTTPAdapter


class OllamaClient:
    """封装Ollama API客户端，支持连接复用和参数化输入"""

    def __init__(self, base_url="http://192.168.50.175:11434", model="llama3.1:8b"):
        self.base_url = base_url
        self.model = model
        self.session = requests.Session()
        # 配置连接池（最大10连接，每主机保持5连接）
        self.session.mount('http://', HTTPAdapter(pool_connections=10, pool_maxsize=5))

    def _build_prompt(self, sample_data, candidate_types):
        """构建标准化提示模板"""
        return f"""
        请根据以下JSON示例值，判断每个字段的数据类型。要求：
        1. 必须从候选类型中选择，不可自创类型 若不清楚类型就选择unknown
        2. 严格按照输出格式输出
        3. 忽略空值和异常值，按多数有效样本判断
        4. 直接输出结果，不要解释
        
        候选类型：{json.dumps(candidate_types, ensure_ascii=False)}
        
        示例数据：
        {json.dumps(sample_data, indent=2, ensure_ascii=False)}
        
        输出格式：
        {{ "字段名1": "类型", "字段名2": "类型", ... }}
                """.strip()

    def analyze_data_type(self, sample_data, candidate_types=None, temperature=0.7, max_tokens=256):
        """执行数据类型分析"""
        if candidate_types is None:
            candidate_types = ["phone", 'address', "email", "date", "id", "username", "int", "city", 'ip', 'unknown']

        prompt = self._build_prompt(sample_data, candidate_types)

        data = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens
            }
        }

        try:
            response = self.session.post(
                f"{self.base_url}/api/generate",
                json=data,
                timeout=60  # 添加超时控制
            )
            response.raise_for_status()
            return response.json().get("response", "No response")
        except requests.exceptions.RequestException as e:
            return f"Error: {str(e)}"
        except json.JSONDecodeError:
            return "Error: Invalid response format"

